package pcsv

import (
	"encoding/csv"
	"log"
	"os"
	"path"
	"strconv"
)

type PCsvWriter struct {
	Lines      int
	FileNo     int
	CsvFile    *os.File
	MaxLines   int
	FilePrefix string
	OutputDir  string
	Title      []string
}

func NewPCsvWriter(filePrefix, outputDir string) *PCsvWriter {
	return &PCsvWriter{
		Lines:      0,
		FileNo:     0,
		CsvFile:    nil,
		MaxLines:   -1,
		FilePrefix: filePrefix,
		OutputDir:  outputDir,
		Title:      nil,
	}
}

func (pw *PCsvWriter) Write(dataCh chan []string) {
	var w = pw.getWriter()
	for data := range dataCh {
		if pw.MaxLines > 0 && pw.Lines >= pw.MaxLines {
			w.Flush()
			w = pw.getWriter()
		}
		if err := w.Write(data); err != nil {
			log.Fatal(err)
		}
		pw.Lines++
	}
	w.Flush()
	pw.CsvFile.Close()
}

func (pw *PCsvWriter) getWriter() *csv.Writer {
	var err error
	if pw.CsvFile != nil {
		pw.CsvFile.Close()
		pw.CsvFile = nil
	}
	outFile := path.Join(pw.OutputDir, pw.FilePrefix)
	if pw.FileNo > 0 {
		outFile += "." + strconv.Itoa(pw.FileNo)
	}

	outFile += ".csv"
	pw.CsvFile, err = os.Create(outFile)
	if err != nil {
		log.Fatalf("pcswriter getWriter err: %v\n", err)
	}
	pw.FileNo++
	pw.Lines = 0
	w := csv.NewWriter(pw.CsvFile)
	// write title
	if pw.Title != nil {
		if err := w.Write(pw.Title); err != nil {
			log.Fatal(err)
		}
	}

	return w
}
